# standard lib
import math
import time
import functools
from itertools import chain

# game imports
import pygame
from pygame.locals import *
from pygame.color import *
import pymunk as pm
from pymunk import Vec2d

# local imports
from data import filepath, load
import constants as c
import sounds

# make albow use the data dir not resources.
import fix
fix.add_parent_dir()
fix.fix_resource()
from albow.resource import get_image
from albow.image_array import get_image_array

X,Y = 0,1
### Physics collision types
COLLTYPE_DEFAULT = 0
COLLTYPE_MOUSE = 1
COLLTYPE_BODY = 2
COLLTYPE_PERL = 3

MAX_TIME = 15

def flipyv(v):
    """Small hack to convert chipmunk physics to pygame coordinates"""
    return (v.x, -v.y+600)

def flipy(y):
    """Small hack to convert chipmunk physics to pygame coordinates"""
    return -y+600

def onscreen(x, y):
    return 0 < x < c.XMAX and 0 < y < c.YMAX

def draw_ball(screen, outer_color, ball):
    r = ball.radius
    v = ball.body.position
    #rot = ball.body.rotation_vector
    p = int(v.x), int(flipy(v.y))
    p2 = int(v.x - r), int(flipy(v.y + r))
    #p2 = Vec2d(rot.x, -rot.y) * r * 0.9
    screen.blit(ball.perl, p2)
    #pygame.draw.circle(screen, THECOLORS[outer_color], p, int(r), 2)
    #pygame.draw.line(screen, THECOLORS["red"], p, p+p2)
    return onscreen(*p)

def draw_line(screen, line, color=None):
    if not color:
        color = "black"

    body = line.body

    pv1 = body.position + line.a.rotated(body.angle)
    pv2 = body.position + line.b.rotated(body.angle)
    p1 = pv1.x, flipy(pv1.y)
    p2 = pv2.x, flipy(pv2.y)
    pygame.draw.lines(screen, THECOLORS[color], False, [p1,p2], 5)

def draw_poly(screen, poly, color=None):
    if not color:
        color = "black"

    body = poly.body
    ps = poly.get_points()
    ps.append(ps[0])
    ps = map(flipyv, ps)
    color = THECOLORS[color]
    transparently = color[:3] + (100,)
    pygame.draw.polygon(screen, transparently, ps, 0)

def draw_bobble(screen, james, bobble, ref):
    rad_angle = bobble.body.angle


    new_james = pygame.transform.rotate(james, math.degrees(rad_angle))
    rect = new_james.get_rect()
    p = ref.position
    rect.center = p.x, flipy(p.y)
    screen.blit(new_james, rect)

def make_ball(scorer, space):
    body = pm.Body(1, 100)
    body.position = (550, 200)
    shape = pm.Circle(body, 10, (0,0))
    shape.friction = 0.1
    shape.elasticity = 0.95
    shape.collision_type = COLLTYPE_PERL
    sprite = pygame.image.load(filepath('perl.png'))
    shape.perl = sprite
    space.add(body, shape)
    scorer.balls.append(shape)
    impulse = Vec2d.unit().rotated(math.radians(53))
    body.apply_impulse(impulse * 1000, (550, 100))
    scorer.numballs += 1
    return time.time()

class BasicScore(object):
    MAX_TIME = 15
    def __init__(self):
        self.numballs = 0
        self.balls = None
        self.time_remaining = self.MAX_TIME
        self.safe_zone = None

    def start(self):
        self._init_time = time.time()

    def is_finished(self):
        now = time.time()
        start_time = self._init_time
        finished = now - start_time > self.MAX_TIME
        self.time_remaining = round(self.MAX_TIME - (now - start_time))
        return finished

    @property
    def saved(self):
        "return number of balls saved."
        contains = self.safe_zone.point_query
        savedballs = filter(lambda b: contains(b.body.position), self.balls)
        return len(savedballs)

    @property
    def total_balls(self):
        return self.numballs



def blit_or_cache(itemname, blit, location):
    try:
        item = getattr(blit_status_bar, itemname)
    except AttributeError:
        item = get_image(itemname + ".png")
        setattr(blit_status_bar, itemname, item)
    blit(item, location)
    right, bottom = item.get_rect().bottomright
    return location[1] + bottom, right + location[0]

def cache_first(fn):
    fn._cache = {}
    def _cached_blit(text, blit, location, colour="black"):
        rendered = fn._cache.get((text, colour))
        if not rendered:
            rendered = fn(text, blit, location, colour)
            fn._cache[(text, colour)] = rendered
        else:
            blit(rendered, location)
        return rendered
    return _cached_blit

@cache_first
def blit_number(text, blit, location, colour):
    image_array = get_image_array("numbers.png", shape=10)
    images = []
    for number in text:
        images.append(image_array[int(number)])

    surf = pygame.Surface((15 * len(images), 20))
    surf.set_colorkey(surf.get_at((0,0)))
    surf.set_alpha(200)
    for i, number in enumerate(images):
        surf.blit(number, (i * 15, 0))

    blit(surf, location)
    return surf

def blit_status_bar(screen, time_remaining, balls_saved, balls_used):
    bottom, right = blit_or_cache("title", screen.blit, (0, 0))
    _, right = blit_or_cache("status_time", screen.blit, (0, bottom))
    blit_number("%02d" % time_remaining, screen.blit, (right + 10, bottom + 5))
    _, right = blit_or_cache("status_lost", screen.blit, (right + 80, bottom))
    lost = (balls_used - balls_saved)
    blit_number("%02d" % lost, screen.blit, (right + 10, bottom + 5))
    _, right = blit_or_cache("status_saved", screen.blit, (right + 80, bottom))
    blit_number("%02d" % balls_saved, screen.blit, (right + 10, bottom + 5))


def main():
    screen = init((800, 600))
    bg = get_image("background.png")
    scorer = level(screen, bg)

def level(screen, bg):
    scorer = BasicScore()
    clock = pygame.time.Clock()
    saved= 0
    for i, surf in enumerate(engine(screen, scorer, bg)):
        if not i % 4:
            saved = scorer.saved

        blit_status_bar(screen, scorer.time_remaining,
                                saved,
                                scorer.total_balls)
        pygame.display.flip()
        clock.tick(50)

    return scorer

def init(size):
    pygame.init()
    screen = pygame.display.set_mode(size)
    return screen

def sound_on_impact(space, arbiter, *args, **kws):
    sounds.play_sfx("hit-james.wav")

def engine(screen, scorer, background=None):

    running = True
    # get james.
    james = pygame.image.load(filepath('james.png'))

    # get bag.
    bag = pygame.image.load(filepath('back-top.png'))

    ### Physics stuff
    pm.init_pymunk()
    space = pm.Space()
    space.gravity = Vec2d(0.0, -900.0)
    space.add_collision_handler(COLLTYPE_PERL, COLLTYPE_BODY, None, None, None,
            sound_on_impact)

    ## Balls
    balls = []
    scorer.balls = balls

    ## Walls
    static = pm.Body(pm.inf, pm.inf)
    walls = [pm.Segment(static, (0, 0), (0, 600), 5.0),
             pm.Segment(static, (800, 0), (800, 600), 5.0),
             pm.Segment(static, (0, 600), (800, 600), 5.0),
             ]
    for wall in walls:
        wall.elasticity = 0.7
    space.add(walls)

    # safe bag.
    bagwalls = [pm.Segment(static, (529, 458), (528, 435), 5.0),
                pm.Segment(static, (528, 435), (542, 409), 5.0),
                pm.Segment(static, (542, 409), (597, 375), 5.0),
                pm.Segment(static, (597, 375), (681, 317), 5.0),
                pm.Segment(static, (681, 317), (751, 305), 5.0),
                pm.Segment(static, (751, 305), (753, 454), 5.0)]
    for wall in bagwalls:
        wall.elasticity = 0.2
    space.add(bagwalls)

    # safe zone.
    safe_body = pm.Body(pm.inf, pm.inf)
    vertices = [(529, 456), (677, 527), (752, 453), (753, 310),
                (684, 313), (537, 415)]
    safe_zone = pm.Poly(safe_body, vertices)
    safe_zone.sensor = 1
    space.add(safe_zone)


    james_center = (300, 200)
    moment = 900000
    ## bobble
    bobble_mass = pm.Body(1000, moment)
    bobble_mass.position = james_center
    bobble_shape = pm.Circle(bobble_mass, 80)
    bobble_shape.elasticity = 0.9
    bobble_shape.collision_type = COLLTYPE_BODY

    # bobble head.
    c_bobble_m = pm.Body(100, moment)
    c_bobble_m.position = james_center[0], james_center[1] + 100
    c_bobble_s = pm.Circle(c_bobble_m, 45)
    c_bobble_s.elasticity = 0.1
    c_bobble_s.collision_type = COLLTYPE_BODY

    # Pin main bobble down.
    anchor = pm.Body(pm.inf, pm.inf)
    anchor.position = bobble_mass.position
    pin = pm.PinJoint(bobble_mass, anchor, (0, 0), (0, 0))
    spring = pm.DampedRotarySpring(bobble_mass, anchor, -0, 80000000, 70000)

    # pin head to body.
    head_pivot = pm.PinJoint(c_bobble_m, bobble_mass, (0, 0), (0, 0))
    c1 = pm.Segment(bobble_mass, (0, 0), (50, 125), 1)
    c2 = pm.Segment(bobble_mass, (0, 0), (-50, 125), 1)
    limit = pm.RotaryLimitJoint(bobble_mass, anchor, -3.14/2, 3.14/2)


    space.add(bobble_mass, bobble_shape, pin, spring, head_pivot, c_bobble_m,
            c_bobble_s, c1, c2, limit)

    scorer.start()
    scorer.safe_zone = safe_zone

    run_physics = True
    cool_down_start = time.time()

    while running:
        running = not scorer.is_finished()

        for event in pygame.event.get():
            if event.type == QUIT:
                running = False
            elif event.type == MOUSEBUTTONDOWN and event.button == 1:
                cool_down_start = make_ball(scorer, space)

            elif event.type == KEYDOWN and event.key == K_f:
                bobble_mass.apply_impulse(Vec2d.unit() * 50000, (100, 0))
            elif event.type == KEYDOWN and event.key == K_j:
                bobble_mass.apply_impulse(Vec2d.unit() * 50000, (-100, 0))

        # hold down mouse for multiball.
        b1, _, _ = pygame.mouse.get_pressed()
        if b1 and time.time() - cool_down_start > c.COOL_DOWN:
            cool_down_start = make_ball(scorer, space)

        ### Update physics
        if run_physics:
            dt = 1.0 / c.FPS
            space.step(dt)

        ### Draw stuff
        if not background:
            screen.fill(THECOLORS["white"])
        else:
            screen.blit(background, (0, 0))

        # Display some text

        bobble_mass.position = anchor.position
        # draw bobble
        draw_bobble(screen, james, bobble_shape, anchor)
        #draw_ball(screen, "purple", bobble_shape)
        #draw_ball(screen, "purple", c_bobble_s)

        #draw_poly(screen, safe_zone, "green")
        for ball in reversed(balls):
            if not draw_ball(screen, "blue", ball):
                balls.remove(ball)

        for line in walls:
            draw_line(screen, line)

        screen.blit(bag, (531, flipy(459)))
        yield screen

if __name__ == '__main__':
    doprof = 0
    if not doprof:
        main()
    else:
        import cProfile, pstats
        prof = cProfile.run("main()", "profile.prof")
        stats = pstats.Stats("profile.prof")
        stats.strip_dirs()
        stats.sort_stats('cumulative', 'time', 'calls')
        stats.print_stats(30)
